package com.ruoyi.gateway.filter;

import com.ruoyi.common.core.utils.StringUtils;
import com.ruoyi.common.core.utils.html.EscapeUtil;
import com.ruoyi.gateway.config.properties.XssProperties;
import io.netty.buffer.ByteBufAllocator;
import java.nio.charset.StandardCharsets;
import javax.annotation.Resource;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * 跨站脚本过滤器
 *
 * @author ruoyi
 */
//@Component
@ConditionalOnProperty(value = "security.xss.enabled", havingValue = "true")
public class XssFilter implements GlobalFilter, Ordered {

  /**
   * 跨站脚本的 xss 配置，nacos自行添加
   */
  @Resource
  private XssProperties xss;

  @Override
  public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
    ServerHttpRequest request = exchange.getRequest();
    // xss开关未开启 或 通过nacos关闭，不过滤
    if (Boolean.FALSE.equals(xss.getEnabled())) {
      return chain.filter(exchange);
    }
    // GET DELETE 不过滤
    HttpMethod method = request.getMethod();
    if (method == null || method == HttpMethod.GET || method == HttpMethod.DELETE) {
      return chain.filter(exchange);
    }
    // 非json类型，不过滤
    if (!isJsonRequest(exchange)) {
      return chain.filter(exchange);
    }
    // excludeUrls 不过滤
    String url = request.getURI().getPath();
    if (StringUtils.matches(url, xss.getExcludeUrls())) {
      return chain.filter(exchange);
    }
    ServerHttpRequestDecorator httpRequestDecorator = requestDecorator(exchange);
    return chain.filter(exchange.mutate().request(httpRequestDecorator).build());

  }

  /**
   * 创建一个ServerHttpRequestDecorator装饰器，用于修改请求体和请求头
   *
   * @param exchange 当前的ServerWebExchange对象，包含请求和响应信息
   * @return 返回一个自定义的ServerHttpRequestDecorator装饰器实例
   */
  private ServerHttpRequestDecorator requestDecorator(ServerWebExchange exchange) {
    // 创建ServerHttpRequestDecorator实例，包装原始请求
    ServerHttpRequestDecorator serverHttpRequestDecorator = new ServerHttpRequestDecorator(exchange.getRequest()) {
      /**
       * 重写getBody方法，用于修改请求体内容
       *
       * @return 返回处理后的Flux<DataBuffer>流
       */
      @Override
      public Flux<DataBuffer> getBody() {
        // 获取原始请求体
        Flux<DataBuffer> body = super.getBody();
        // 缓冲请求体数据，并对每个缓冲区进行处理
        return body.buffer().map(dataBuffers -> {
          // 创建默认数据缓冲区工厂
          DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
          // 合并所有数据缓冲区
          DataBuffer join = dataBufferFactory.join(dataBuffers);
          // 创建字节数组存储请求内容
          byte[] content = new byte[join.readableByteCount()];
          // 读取请求内容到字节数组
          join.read(content);
          // 释放数据缓冲区
          DataBufferUtils.release(join);
          // 将字节数组转换为字符串
          String bodyStr = new String(content, StandardCharsets.UTF_8);
          // 防xss攻击过滤
          bodyStr = EscapeUtil.clean(bodyStr);
          // 转成字节
          byte[] bytes = bodyStr.getBytes(StandardCharsets.UTF_8);
          NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(ByteBufAllocator.DEFAULT);
          DataBuffer buffer = nettyDataBufferFactory.allocateBuffer(bytes.length);
          buffer.write(bytes);
          return buffer;
        });
      }

      @Override
      public HttpHeaders getHeaders() {
        HttpHeaders httpHeaders = new HttpHeaders();
        httpHeaders.putAll(super.getHeaders());
        // 由于修改了请求体的body，导致content-length长度不确定，因此需要删除原先的content-length
        httpHeaders.remove(HttpHeaders.CONTENT_LENGTH);
        httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
        return httpHeaders;
      }

    };
    return serverHttpRequestDecorator;
  }

  /**
   * 是否是Json请求
   *
   * @param exchange HTTP请求
   */
  public boolean isJsonRequest(ServerWebExchange exchange) {
    String header = exchange.getRequest().getHeaders().getFirst(HttpHeaders.CONTENT_TYPE);
    return StringUtils.startsWithIgnoreCase(header, MediaType.APPLICATION_JSON_VALUE);
  }

  @Override
  public int getOrder() {
    return -100;
  }
}
